# ---------------------------------------------
# Copyright (c) OpenMMLab. All rights reserved.
# ---------------------------------------------
#  Modified by Zhiqi Li
# ---------------------------------------------

import torch
from mmcv.runner import force_fp32, auto_fp16
from mmdet.models import DETECTORS
from mmdet3d.core import bbox3d2result
from mmdet3d.models.detectors.mvx_two_stage import MVXTwoStageDetector
from projects.mmdet3d_plugin.models.utils.grid_mask import GridMask
import time
import copy
import numpy as np
import mmdet3d
from projects.mmdet3d_plugin.models.utils.bricks import run_time


@DETECTORS.register_module()
class BEVFormer(MVXTwoStageDetector):
    """BEVFormer.
    Args:
        video_test_mode (bool): Decide whether to use temporal information during inference.
    """

    def __init__(self,
                 use_grid_mask=False,
                 pts_voxel_layer=None,
                 pts_voxel_encoder=None,
                 pts_middle_encoder=None,
                 pts_fusion_layer=None,
                 img_backbone=None,
                 pts_backbone=None,
                 img_neck=None,
                 pts_neck=None,
                 pts_bbox_head=None,
                 img_roi_head=None,
                 img_rpn_head=None,
                 train_cfg=None,
                 test_cfg=None,
                 pretrained=None,
                 video_test_mode=False
                 ):

        super(BEVFormer,
              self).__init__(pts_voxel_layer, pts_voxel_encoder,
                             pts_middle_encoder, pts_fusion_layer,
                             img_backbone, pts_backbone, img_neck, pts_neck,
                             pts_bbox_head, img_roi_head, img_rpn_head,
                             train_cfg, test_cfg, pretrained)
        self.grid_mask = GridMask(
            True, True, rotate=1, offset=False, ratio=0.5, mode=1, prob=0.7)
        self.use_grid_mask = use_grid_mask
        self.fp16_enabled = False

        # temporal
        self.video_test_mode = video_test_mode
        self.prev_frame_info = {
            'prev_bev': None,
            'scene_token': None,
            'prev_pos': 0,
            'prev_angle': 0,
        }


    def extract_img_feat(self, img, img_metas, len_queue=None):
        """Extract features of images.

        Args:
            img (torch.Tensor): 输入图像, 维度
                - 时序模式: [bs * len_queue, num_cams, C, H, W]
                - 非时序模式: [bs, num_cams, C, H, W]
            img_metas (list): 图像元数据
            len_queue (int, optional): 时序队列长度

        Returns:
            img_feats_reshaped (list[Tensor]): 每层特征, 维度
                - 时序模式: [bs, len_queue, num_cams, C, H_feat, W_feat]
                - 非时序模式: [bs, num_cams, C, H_feat, W_feat]
        """
        B = img.size(0)  # B: 批大小 (时序模式=bs*len_queue, 非时序模式=bs)
        if img is not None:
            # input_shape = img.shape[-2:]
            # # update real input shape of each single img
            # for img_meta in img_metas:
            #     img_meta.update(input_shape=input_shape)

            # img: [B, N, C, H, W]
            if img.dim() == 5 and img.size(0) == 1:
                # squeeze: [1, N, C, H, W] -> [N, C, H, W]
                img.squeeze_()
            elif img.dim() == 5 and img.size(0) > 1:
                B, N, C, H, W = img.size()
                # [B, N, C, H, W] -> [B*N, C, H, W]
                img = img.reshape(B * N, C, H, W)
            # img: [B*N, C, H, W]

            """
                叠加网格状掩码，掩盖部分图像
                好处：
                    提高鲁棒性: 强制模型学会从部分信息中进行推理
                    防止过拟合: 随机掩码防止模型过度依赖特定区域
                    模拟遮挡: 模拟真实场景中的遮挡情况
                    增强泛化: 提高模型在不同条件下的性能
            """
            if self.use_grid_mask:
                img = self.grid_mask(img)  # 维度不变

            img_feats = self.img_backbone(img)
            # img_feats: List[Tensor], 每层特征
            # img_feats[i]: [BN, C, H_feat, W_feat]
            # BN = B*N 或 len_queue*num_cams
            if isinstance(img_feats, dict):
                img_feats = list(img_feats.values())
        else:
            return None
        if self.with_img_neck:
            img_feats = self.img_neck(img_feats)
            # img_feats: List[Tensor], 每层特征
            # img_feats[i]: [BN, C, H_feat, W_feat]

        img_feats_reshaped = []
        for img_feat in img_feats:
            BN, C, H, W = img_feat.size()
            # BN = B*N 或 len_queue*num_cams
            if len_queue is not None:
                # 时序模式: [BN, C, H, W] -> [bs, len_queue, num_cams, C, H, W]
                img_feats_reshaped.append(
                    img_feat.view(int(B/len_queue), len_queue, int(BN / B), C, H, W)
                )
                # int(B/len_queue): bs
                # len_queue: 时序长度
                # int(BN/B): num_cams
            else:
                # 非时序模式: [BN, C, H, W] -> [B, num_cams, C, H, W]
                img_feats_reshaped.append(
                    img_feat.view(B, int(BN / B), C, H, W)
                )
                # B: 批大小
                # int(BN/B): num_cams
        # img_feats_reshaped: List[Tensor], 每层特征
        # 时序: [bs, len_queue, num_cams, C, H, W]
        # 非时序: [bs, num_cams, C, H, W]
        return img_feats_reshaped

    @auto_fp16(apply_to=('img'))
    def extract_feat(self, img, img_metas=None, len_queue=None):
        """Extract features from images and points."""

        img_feats = self.extract_img_feat(img, img_metas, len_queue=len_queue)
        
        return img_feats


    def forward_pts_train(self,
                          pts_feats,
                          gt_bboxes_3d,
                          gt_labels_3d,
                          img_metas,
                          gt_bboxes_ignore=None,
                          prev_bev=None):
        """Forward function'
        Args:
            pts_feats (list[torch.Tensor]): Features of point cloud branch
            gt_bboxes_3d (list[:obj:`BaseInstance3DBoxes`]): Ground truth
                boxes for each sample.
            gt_labels_3d (list[torch.Tensor]): Ground truth labels for
                boxes of each sampole
            img_metas (list[dict]): Meta information of samples.
            gt_bboxes_ignore (list[torch.Tensor], optional): Ground truth
                boxes to be ignored. Defaults to None.
            prev_bev (torch.Tensor, optional): BEV features of previous frame.
        Returns:
            dict: Losses of each branch.
        """

        outs = self.pts_bbox_head(
            pts_feats, img_metas, prev_bev)
        loss_inputs = [gt_bboxes_3d, gt_labels_3d, outs]
        losses = self.pts_bbox_head.loss(*loss_inputs, img_metas=img_metas)
        return losses

    def forward_dummy(self, img):
        dummy_metas = None
        return self.forward_test(img=img, img_metas=[[dummy_metas]])

    def forward(self, return_loss=True, **kwargs):
        """Calls either forward_train or forward_test depending on whether
        return_loss=True.
        Note this setting will change the expected inputs. When
        `return_loss=True`, img and img_metas are single-nested (i.e.
        torch.Tensor and list[dict]), and when `resturn_loss=False`, img and
        img_metas should be double nested (i.e.  list[torch.Tensor],
        list[list[dict]]), with the outer list indicating test time
        augmentations.
        """
        if return_loss:
            return self.forward_train(**kwargs)
        else:
            return self.forward_test(**kwargs)
    
    def obtain_history_bev(self, imgs_queue, img_metas_list):
        """Obtain history BEV features iteratively. To save GPU memory, gradients are not calculated.

        Args:
            imgs_queue: torch.Tensor, 历史图像队列, 维度 [bs, len_queue, num_cams, C, H, W]
                - bs: batch size (通常为1)
                - len_queue: 历史帧数 (如2)
                - num_cams: 摄像头数量 (如6)
                - C: 通道数 (3)
                - H, W: 图像高宽 (480, 800)
            img_metas_list: List[List[Dict]], 历史帧元数据, 维度 [bs][len_queue]
                - 每个dict为一帧的元数据

        Returns:
            prev_bev: torch.Tensor, 历史BEV特征, 维度 [bs, bev_h*bev_w, embed_dims] (如[1, 2500, 256])
        """
        self.eval()

        with torch.no_grad():
            prev_bev = None
            bs, len_queue, num_cams, C, H, W = imgs_queue.shape  # [1, 2, 6, 3, 480, 800]
            imgs_queue = imgs_queue.reshape(bs*len_queue, num_cams, C, H, W)  # [2, 6, 3, 480, 800]
            img_feats_list = self.extract_feat(img=imgs_queue, len_queue=len_queue)
            # img_feats_list: List[Tensor], 每层特征, 维度 [bs*len_queue, num_cams, embed_dims, H_feat, W_feat]
            for i in range(len_queue):
                img_metas = [each[i] for each in img_metas_list]  # 当前帧元数据, List长度为bs
                if not img_metas[0]['prev_bev_exists']:
                    prev_bev = None
                # img_feats: List[Tensor], 每层特征, 维度 [bs, num_cams, embed_dims, H_feat, W_feat]
                img_feats = [each_scale[:, i] for each_scale in img_feats_list]
                # each_scale: [bs, len_queue, num_cams, embed_dims, H_feat, W_feat]
                # each_scale[:, i]: [bs, num_cams, embed_dims, H_feat, W_feat]
                # 通过配置可以看出跳转到pts_bbox_head的forward方法
                prev_bev = self.pts_bbox_head(
                    img_feats, img_metas, prev_bev, only_bev=True)
                # prev_bev: torch.Tensor, [bs, bev_h*bev_w, embed_dims]
            self.train()
            return prev_bev

    @auto_fp16(apply_to=('img', 'points'))
    def forward_train(self,
                      points=None,
                      img_metas=None,
                      gt_bboxes_3d=None,
                      gt_labels_3d=None,
                      gt_labels=None,
                      gt_bboxes=None,
                      img=None,
                      proposals=None,
                      gt_bboxes_ignore=None,
                      img_depth=None,
                      img_mask=None,
                      ):
        """Forward training function.
        Args:
            points (list[torch.Tensor], optional): Points of each sample.
                Defaults to None.
            img_metas (list[dict], optional): Meta information of each sample.
                Defaults to None.
            gt_bboxes_3d (list[:obj:`BaseInstance3DBoxes`], optional):
                Ground truth 3D boxes. Defaults to None.
            gt_labels_3d (list[torch.Tensor], optional): Ground truth labels
                of 3D boxes. Defaults to None.
            gt_labels (list[torch.Tensor], optional): Ground truth labels
                of 2D boxes in images. Defaults to None.
            gt_bboxes (list[torch.Tensor], optional): Ground truth 2D boxes in
                images. Defaults to None.
            img (torch.Tensor optional): Images of each sample with shape
                (N, C, H, W). Defaults to None.
            proposals ([list[torch.Tensor], optional): Predicted proposals
                used for training Fast RCNN. Defaults to None.
            gt_bboxes_ignore (list[torch.Tensor], optional): Ground truth
                2D boxes in images to be ignored. Defaults to None.
        Returns:
            dict: Losses of different branches.
        """
        # ==== 主要维度总结 ====
        # 1. 输入图像: [bs, len_queue, num_cams, 3, H, W] = [1, 3, 6, 3, 480, 800]
        # 2. 历史图像: [bs, len_queue-1, num_cams, 3, H, W] = [1, 2, 6, 3, 480, 800]  
        # 3. 当前图像: [bs, num_cams, 3, H, W] = [1, 6, 3, 480, 800]
        # 4. 图像特征: List[[bs, num_cams, embed_dims, H_feat, W_feat]] = [[1, 6, 256, 15, 25]]
        # 5. 历史BEV: [bs, bev_h*bev_w, embed_dims] = [1, 2500, 256]
        # 6. GT 3D框: List[BaseInstance3DBoxes] 每个[num_gt, 9]
        # 7. GT标签: List[torch.Tensor] 每个[num_gt]
        
        # ==== 输入变量维度分析 ====
        # img: [bs, len_queue, num_cams, C, H, W] = [1, 3, 6, 3, 480, 800]
        # - bs: batch size (通常为1)
        # - len_queue: 时序队列长度 (3帧历史 + 1帧当前)
        # - num_cams: 摄像头数量 (nuScenes有6个摄像头)
        # - C: 图像通道数 (RGB=3)
        # - H, W: 图像高度和宽度 (480x800 经过resize后)
        
        # img_metas: List[List[Dict]] 长度为 [bs][len_queue]
        # - 每个dict包含摄像头标定参数、变换矩阵等元数据
        
        # gt_bboxes_3d: List[BaseInstance3DBoxes] 长度为 bs
        # - 每个元素包含该帧的所有3D真值框 [num_gt, 9] (x,y,z,w,l,h,vx,vy,yaw)
        
        # gt_labels_3d: List[torch.Tensor] 长度为 bs  
        # - 每个tensor维度 [num_gt] 包含3D框的类别标签 (0-9, 对应10个类别)
        
        len_queue = img.size(1) # 3 (从img维度获取时序长度)
        
        # ==== 时序数据分离 ====
        prev_img = img[:, :-1, ...]  # [1, 2, 6, 3, 480, 800] - 历史帧 (前len_queue-1帧)
        img = img[:, -1, ...]        # [1, 6, 3, 480, 800] - 当前帧 (最后1帧)

        # ==== 元数据处理 ====
        '''图像元数据结构说明:
        img_metas = [  # List长度为batch_size
            [  # List长度为len_queue，每帧一个dict
                {
                    'filename': ['cam0.jpg', 'cam1.jpg', ..., 'cam5.jpg'],  # 6个摄像头图片路径
                    'img_shape': (480, 800, 3),                            # 图像原始尺寸
                    'pad_shape': (480, 832, 3),                            # padding后尺寸  
                    'scale_factor': 1.0,                                   # 图像缩放比例
                    'lidar2img': [np.array(4x4), ..., np.array(4x4)],      # 6个LiDAR到图像变换矩阵
                    'intrinsics': [np.array(3x3), ..., np.array(3x3)],     # 6个摄像头内参矩阵
                    'extrinsics': [np.array(4x4), ..., np.array(4x4)],     # 6个摄像头外参矩阵
                    'prev_bev_exists': True,                               # 是否使用历史BEV特征
                    'can_bus': np.array(18),                               # 车辆状态信息(位置、角度等)
                    'flip': False,                                         # 数据增强-翻转标志
                }, ...
            ]
        ]
        '''
        prev_img_metas = copy.deepcopy(img_metas)  # 深拷贝历史帧元数据
        
        # ==== 历史BEV特征获取 ====
        # obtain_history_bev函数输入输出维度:
        # 输入: prev_img [1, 2, 6, 3, 480, 800], prev_img_metas
        # 输出: prev_bev [1, bev_h*bev_w, embed_dims] = [1, 2500, 256] (对于tiny版本)
        #       其中 bev_h=bev_w=50, embed_dims=256
        prev_bev = self.obtain_history_bev(prev_img, prev_img_metas)

        # ==== 当前帧元数据提取 ====
        img_metas = [each[len_queue-1] for each in img_metas]  # 提取最后一帧的元数据
        if not img_metas[0]['prev_bev_exists']:
            prev_bev = None
            
        # ==== 当前帧特征提取 ====
        # extract_feat函数输入输出维度:
        # 输入: img [1, 6, 3, 480, 800]
        # 经过backbone(ResNet50)+neck(FPN)后:
        # 输出: img_feats List[torch.Tensor] 长度为num_levels (tiny版本=1)
        #       img_feats[0]: [1, 6, 256, H_feat, W_feat]
        #       其中H_feat, W_feat取决于backbone的下采样倍数
        #       对于ResNet50 stage4: stride=32, 所以 H_feat=480//32=15, W_feat=800//32=25
        #       最终: img_feats[0] = [1, 6, 256, 15, 25]
        img_feats = self.extract_feat(img=img, img_metas=img_metas)
        
        # ==== 损失计算 ====
        losses = dict()  # 初始化损失字典
        
        # forward_pts_train函数输入输出维度:
        # 输入:
        # - img_feats: List[[1, 6, 256, 15, 25]] 
        # - gt_bboxes_3d: List[BaseInstance3DBoxes] 包含真值3D框
        # - gt_labels_3d: List[torch.Tensor] 包含真值标签 [num_gt]
        # - img_metas: List[Dict] 当前帧元数据
        # - prev_bev: [1, 2500, 256] 历史BEV特征
        # 输出: losses dict 包含分类损失、回归损失等
        losses_pts = self.forward_pts_train(img_feats, gt_bboxes_3d,
                                            gt_labels_3d, img_metas,
                                            gt_bboxes_ignore, prev_bev)

        losses.update(losses_pts)
        return losses

    def forward_test(self, img_metas, img=None, **kwargs):
        for var, name in [(img_metas, 'img_metas')]:
            if not isinstance(var, list):
                raise TypeError('{} must be a list, but got {}'.format(
                    name, type(var)))
        img = [img] if img is None else img

        if img_metas[0][0]['scene_token'] != self.prev_frame_info['scene_token']:
            # the first sample of each scene is truncated
            self.prev_frame_info['prev_bev'] = None
        # update idx
        self.prev_frame_info['scene_token'] = img_metas[0][0]['scene_token']

        # do not use temporal information
        if not self.video_test_mode:
            self.prev_frame_info['prev_bev'] = None

        # Get the delta of ego position and angle between two timestamps.
        tmp_pos = copy.deepcopy(img_metas[0][0]['can_bus'][:3])
        tmp_angle = copy.deepcopy(img_metas[0][0]['can_bus'][-1])
        if self.prev_frame_info['prev_bev'] is not None:
            img_metas[0][0]['can_bus'][:3] -= self.prev_frame_info['prev_pos']
            img_metas[0][0]['can_bus'][-1] -= self.prev_frame_info['prev_angle']
        else:
            img_metas[0][0]['can_bus'][-1] = 0
            img_metas[0][0]['can_bus'][:3] = 0

        new_prev_bev, bbox_results = self.simple_test(
            img_metas[0], img[0], prev_bev=self.prev_frame_info['prev_bev'], **kwargs)
        # During inference, we save the BEV features and ego motion of each timestamp.
        self.prev_frame_info['prev_pos'] = tmp_pos
        self.prev_frame_info['prev_angle'] = tmp_angle
        self.prev_frame_info['prev_bev'] = new_prev_bev
        return bbox_results

    def simple_test_pts(self, x, img_metas, prev_bev=None, rescale=False):
        """Test function"""
        outs = self.pts_bbox_head(x, img_metas, prev_bev=prev_bev)

        bbox_list = self.pts_bbox_head.get_bboxes(
            outs, img_metas, rescale=rescale)
        bbox_results = [
            bbox3d2result(bboxes, scores, labels)
            for bboxes, scores, labels in bbox_list
        ]
        return outs['bev_embed'], bbox_results

    def simple_test(self, img_metas, img=None, prev_bev=None, rescale=False):
        """Test function without augmentaiton."""
        img_feats = self.extract_feat(img=img, img_metas=img_metas)

        bbox_list = [dict() for i in range(len(img_metas))]
        new_prev_bev, bbox_pts = self.simple_test_pts(
            img_feats, img_metas, prev_bev, rescale=rescale)
        for result_dict, pts_bbox in zip(bbox_list, bbox_pts):
            result_dict['pts_bbox'] = pts_bbox
        return new_prev_bev, bbox_list
